package com.zelix.analy;

import com.arthas.client.ArthasClient;
import com.arthas.client.ArthasTool;
import me.coley.analysis.SimAnalyzer;
import me.coley.analysis.SimInterpreter;
import me.coley.analysis.value.AbstractValue;
import org.objectweb.asm.ClassReader;
import org.objectweb.asm.Opcodes;
import org.objectweb.asm.Type;
import org.objectweb.asm.tree.*;
import org.objectweb.asm.tree.analysis.AnalyzerException;
import org.objectweb.asm.tree.analysis.Frame;

import java.io.IOException;
import java.util.Enumeration;
import java.util.HashMap;
import java.util.List;
import java.util.jar.JarEntry;
import java.util.jar.JarFile;
import java.util.zip.ZipEntry;

import static com.zelix.analy.TTT.getKeys;

public class Flow {

    static HashMap<String, String> keyMap = getKeys("./file/lesslong.txt");
    static HashMap<String, String> skips = getKeys("./file/skip.txt");

    public static void main(String[] args) throws IOException {
        JarFile jarFile = null;
        try {
            jarFile = new JarFile("./file/ZKM_bak.jar");
        } catch (IOException e) {
            e.printStackTrace();
        }

        ArthasClient client = ArthasTool.getClient();

        if (jarFile == null || client == null) {
            return;
        }

        String rule = "com/zelix/jr";

        Enumeration<JarEntry> entries = jarFile.entries();
        while (entries.hasMoreElements()) {
            JarEntry entry = entries.nextElement();
            if (entry.getName().equals(rule + ".class")) {//只分析class文件
                //去除文件名后缀 ".class"
                analy(client, jarFile, rule);
            }
        }

        client.disconnect();
    }

    public static void analy(ArthasClient client, JarFile jarFile, String cls) throws IOException {
        System.out.println("***********" + cls + "***********");

        ZipEntry entry = jarFile.getEntry(cls + ".class");

        if (entry != null) {
            ClassReader cr = new ClassReader(jarFile.getInputStream(entry));
            ClassNode cn = new ClassNode();
            cr.accept(cn, ClassReader.EXPAND_FRAMES);

            List<MethodNode> methods = cn.methods;

            for (MethodNode mn : methods) {
                System.out.println("------------" + mn.name + mn.desc + "------------");
                //计算long值,遍历每个方法前初始化,不同方法值不同
                String key = keyMap.get(cls + ":" + mn.name + ":" + mn.desc);
                if (key != null) {

                    //插入关键key到指令中
                    insertKeyValue(cls, mn);

                    if ((mn.name + mn.desc).startsWith("-O([Ljava/lang/Object;)V")) {
                        System.out.println("error , skip");
                        continue;
                    }

                    //插入后获取frames
                    Frame<AbstractValue>[] frames = getFrames(cls, mn);

                    if (frames == null) {
                        System.out.println("frames null");
                        continue;
                    }

                    //通过frames获取调用其他方法的long参数
                    HashMap<Integer, String> keys = getMethodLongArgs(cls, mn, frames);
                    if (keys.isEmpty()) {
                        System.out.println("keys empty");
                        continue;
                    }
                    System.out.println(keys);

                    InsnList insnList = mn.instructions;

                    //遍历INVOKEDYNAMIC命令
                    for (int index = 0; index < insnList.size(); index++) {
                        if (insnList.get(index) instanceof InvokeDynamicInsnNode) {
                            InvokeDynamicInsnNode dynamicNode = (InvokeDynamicInsnNode) insnList.get(index);
                            if (dynamicNode.desc.equals("(IJ)Ljava/lang/String;") && dynamicNode.bsm.getOwner().equals(cls)) {
                                //字符串解密,无需处理,跳过
                                continue;
                            } else if (dynamicNode.desc.contains("JJ)")) {
                                //System.out.println("方法或者属性调用:");
                                //System.out.println(dynamicNode.name + ":" + dynamicNode.desc + ":" + Arrays.deepToString(dynamicNode.bsmArgs));
                                //System.out.println(dynamicNode.bsm.getTag() + ":" + dynamicNode.bsm.getOwner() + ":" + dynamicNode.bsm.getName() + ":" + dynamicNode.bsm.getDesc());

                                Frame<AbstractValue> frame = frames[index];
                                if (frame != null && frame.getStackSize() > 2) {//参数超过两个long值,INVOKEDYNAMIC指令的两个long
                                    AbstractInsnNode node = singleCmd.getFieldAndMethodNode(client, dynamicNode.name, frame.getStack(frame.getStackSize() - 2).getValue().toString(), frame.getStack(frame.getStackSize() - 1).getValue().toString());
                                    if (node != null && node instanceof MethodInsnNode) {    //只分析MethodInsnNode方法
                                        MethodInsnNode methodInsnNode = (MethodInsnNode) node;
                                        if (methodInsnNode.owner.startsWith("com/zelix/")) {//只分析内部调用

                                            List<AbstractInsnNode> insns = frame.getStack(frame.getStackSize() - 3).getInsns();
                                            for (int i = 0; i < insns.size(); i++) {
                                                //参数为Object[]类型
                                                if (insns.get(i).getOpcode() == Opcodes.ANEWARRAY) {
                                                    //获取这个指令在方法中的位置
                                                    int pos = insnList.indexOf(insns.get(i));
                                                    //得到这个指令的frame
                                                    Frame<AbstractValue> valueFrame = frames[pos];

                                                    if (valueFrame.getStack(valueFrame.getStackSize() - 1).getValue() == null) {
                                                        continue;
                                                    }
                                                    //ANEWARRAY参数栈上一个参数是数组大小
                                                    int arrSize = Integer.valueOf(valueFrame.getStack(valueFrame.getStackSize() - 1).getValue().toString());
                                                    if (arrSize > 0) {//数组不为空时
                                                        //valueFrame.getStackSize() - 1 为数组大小,数组倒序从valueFrame.getStackSize() - 2开始
                                                        //倒序最接近数组,参数为数组参数
                                                        for (int j = valueFrame.getStackSize() - 2; j >= 0; j--) {//倒序查找总的参数栈
                                                            if (valueFrame.getStack(j).getType() == Type.LONG_TYPE) {
                                                                //valueFrame.getStackSize() - 2 -j 倒序索引,从0开始
                                                                //arrSize-1-(valueFrame.getStackSize() - 2 -j)计算正序索引
                                                                //简化为arrSize + 1 +j-valueFrame.getStackSize()
                                                                int keyIndex = arrSize + 1 + j - valueFrame.getStackSize();

                                                                if (valueFrame.getStack(j).getValue() == null) {
                                                                    continue;
                                                                }
                                                                String value = valueFrame.getStack(j).getValue().toString();
                                                                //是否是long参数里的
                                                                if (keys.values().contains(value)) {
                                                                    String rtn = analyFlow.analy(client, jarFile, methodInsnNode.owner, methodInsnNode.name, methodInsnNode.desc, keyIndex, value);
                                                                    //不存在的适合打印
                                                                    if (!skips.containsKey(methodInsnNode.owner + ":" + methodInsnNode.name + ":" + methodInsnNode.desc) && !keyMap.containsKey(methodInsnNode.owner + ":" + methodInsnNode.name + ":" + methodInsnNode.desc)) {
                                                                        if (rtn != null) {
                                                                            System.out.println(rtn);
                                                                        } else {
                                                                            System.out.println("未计算出:" + methodInsnNode.owner + ":" + methodInsnNode.name + ":" + methodInsnNode.desc);
                                                                            System.out.println("数组大小:" + arrSize);
                                                                            System.out.println("key索引:" + keyIndex + "\t值:" + value);
                                                                        }
                                                                    }
                                                                }
                                                            }
                                                        }
                                                    }
                                                }
                                            }
                                        }
                                    }
                                }

                            } else if (dynamicNode.bsm.getOwner().equals("java/lang/invoke/LambdaMetafactory")) {
                                System.out.println("Lambda表达式:" + dynamicNode.name + ":" + dynamicNode.desc);
                            } else {
                                System.out.println("未知用途:");
                                System.out.println("实际方法描述:" + dynamicNode.name + ":" + dynamicNode.desc);
                                System.out.println("调用方法描述:" + dynamicNode.bsm.getOwner() + ":" + dynamicNode.bsm.getName());
                            }
                        } else if (insnList.get(index) instanceof MethodInsnNode) {
                            MethodInsnNode methodInsnNode = (MethodInsnNode) insnList.get(index);
                            if (methodInsnNode.owner.startsWith("com/zelix/") && methodInsnNode.desc.contains("J")) {//只分析内部调用,且参数有long
                                Frame<AbstractValue> frame = frames[index];
                                if (frame != null && frame.getStackSize() > 1) {//参数超过两个long值
                                    for (int i = 0; i < frame.getStackSize(); i++) {
                                        if (frame.getStack(i).getType() == Type.LONG_TYPE) {
                                            if (frame.getStack(i).getValue() == null) {
                                                continue;
                                            }
                                            String value = frame.getStack(i).getValue().toString();
                                            //是否是long参数里的
                                            if (keys.values().contains(value)) {
                                                String rtn = analyFlow.analy(client, jarFile, methodInsnNode.owner, methodInsnNode.name, methodInsnNode.desc, getKeysIndex(keys, value), value);
                                                //不存在的适合打印
                                                if (!skips.containsKey(methodInsnNode.owner + ":" + methodInsnNode.name + ":" + methodInsnNode.desc) && !keyMap.containsKey(methodInsnNode.owner + ":" + methodInsnNode.name + ":" + methodInsnNode.desc)) {
                                                    if (rtn != null) {
                                                        System.out.println(rtn);
                                                    } else {
                                                        System.out.println("未计算出:" + methodInsnNode.owner + ":" + methodInsnNode.name + ":" + methodInsnNode.desc);
                                                        System.out.println("varIndex:" + getKeysIndex(keys, value) + "\t值:" + value);
                                                    }
                                                }
                                            }
                                        }
                                    }
                                }
                            }
                        }
                    }
                }
            }
        }
    }

    /**
     * 插入关键值key
     *
     * @param cls
     * @param mn
     */
    public static void insertKeyValue(String cls, MethodNode mn) {
        InsnList insnList = mn.instructions;

        //计算long值,遍历每个方法前初始化,不同方法值不同
        String key = keyMap.get(cls + ":" + mn.name + ":" + mn.desc);

        //获取关键值指令索引
        int keyIndex = getKeyIndex(insnList);

        if (key != null && keyIndex != -1) {
            //System.out.println(keyIndex + ":" + key);

            //key的指令
            VarInsnNode varInsnNode = (VarInsnNode) insnList.get(keyIndex);

            //插入新指令,把关键key插入,方便frame进行计算
            InsnList tmp = new InsnList();
            tmp.add(new InsnNode(Opcodes.POP2));//弹出
            tmp.add(new LdcInsnNode(Long.valueOf(key)));//设置
            insnList.insertBefore(varInsnNode, tmp);
        }
    }

    /**
     * 获取Frames
     *
     * @param cls
     * @param mn
     * @return
     */
    public static Frame<AbstractValue>[] getFrames(String cls, MethodNode mn) {
        //frame分析
        SimAnalyzer analyzer = new SimAnalyzer(new SimInterpreter());
        analyzer.setThrowUnresolvedAnalyzerErrors(false);
        Frame<AbstractValue>[] frames = null;
        try {
            frames = analyzer.analyze(cls, mn);
        } catch (AnalyzerException e) {
            //e.printStackTrace();
            return null;
        }
        return frames;
    }

    /**
     * 获取方法内调用其他方法的long参数
     *
     * @param cls
     * @param mn
     * @return
     */
    public static HashMap<Integer, String> getMethodLongArgs(String cls, MethodNode mn, Frame<AbstractValue>[] frames) {
        HashMap<Integer, String> keys = new HashMap<>();

        if (frames == null) {
            return keys;
        }

        InsnList insnList = mn.instructions;
        for (int index = 0; index < insnList.size(); index++) {
            if (insnList.get(index) instanceof VarInsnNode && insnList.get(index).getOpcode() == Opcodes.LSTORE) {
                VarInsnNode varInsnNode = (VarInsnNode) insnList.get(index);

                Frame<AbstractValue> frame = frames[index];
                if (frame != null && frame.getStackSize() > 0) {
                    //跳过函数入参等导致计算值为null的情况
                    if (frame.getStack(frame.getStackSize() - 1).getValue() == null) {
                        continue;
                    }
                    keys.put(varInsnNode.var, frame.getStack(frame.getStackSize() - 1).getValue().toString());
                }
            }
        }

        return keys;
    }

    /**
     * 获取key的指令索引
     *
     * @param insnList
     * @return
     */
    public static int getKeyIndex(InsnList insnList) {
        //需要倒序查找最后一个保存命令,可能保存很多次,变量代号都是一个
        for (int index = insnList.size() - 1; index >= 0; index--) {
            //查找LSTORE指令
            if (insnList.get(index) instanceof VarInsnNode && insnList.get(index).getOpcode() == Opcodes.LSTORE) {
                VarInsnNode varInsnNode = (VarInsnNode) insnList.get(index);
                if (isDynamicKey(insnList, varInsnNode.var)) {
                    return index;
                }
            }
        }
        return -1;
    }


    /**
     * 确认这个key是不是关键key
     *
     * @param insnList
     * @param keyVar
     * @return
     */
    public static boolean isDynamicKey(InsnList insnList, int keyVar) {
        for (int index = 0; index < insnList.size(); index++) {
            //查找INVOKEDYNAMIC指令
            if (insnList.get(index) instanceof InvokeDynamicInsnNode) {
                InvokeDynamicInsnNode dynamicNode = (InvokeDynamicInsnNode) insnList.get(index);
                //System.out.println(dynamicNode.name + ":" + dynamicNode.desc + ":" + Arrays.deepToString(dynamicNode.bsmArgs));
                //System.out.println(dynamicNode.bsm.getTag()+ ":" +dynamicNode.bsm.getOwner()+ ":" +dynamicNode.bsm.getName()+ ":" +dynamicNode.bsm.getDesc());
                /**
                 * SIPUSH 19546
                 * LDC 1323103721956097495L
                 * LLOAD 1
                 * LXOR
                 * INVOKEDYNAMIC j (IJ)Ljava/lang/String; handle[H_INVOKESTATIC com/zelix/ZKM.b(Ljava/lang/invoke/MethodHandles$Lookup;Ljava/lang/String;Ljava/lang/invoke/MethodType;)Ljava/lang/invoke/CallSite;] args[]
                 */
                /**
                 * LDC -5040441400421082001L
                 * LLOAD 1
                 * INVOKEDYNAMIC k (Ljava/lang/Object;JJJ)Ljava/lang/StringBuilder; handle[H_INVOKESTATIC com/zelix/v99.a(Ljava/lang/invoke/MethodHandles$Lookup;Ljava/lang/String;Ljava/lang/invoke/MethodType;)Ljava/lang/invoke/CallSite;] args[]
                 */
                if (insnList.get(index - 1) instanceof VarInsnNode || (insnList.get(index - 1) instanceof InsnNode && insnList.get(index - 1).getOpcode() == Opcodes.LXOR && insnList.get(index - 2) instanceof VarInsnNode)) {
                    VarInsnNode varInsnNode = (VarInsnNode) (insnList.get(index - 1) instanceof VarInsnNode ? insnList.get(index - 1) : insnList.get(index - 2));
                    if (keyVar == varInsnNode.var) {
                        return true;
                    }
                }
            }
        }
        return false;
    }

    public static int getKeysIndex(HashMap<Integer, String> keys, String value) {
        for (int i : keys.keySet()) {
            if (keys.get(i).equals(value)) {
                return i;
            }
        }
        return -1;
    }
}
